from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
from openpyxl import Workbook
from db_Session import get_db_session


def export_to_xlsx(table_name, output_file):
    try:
        # 获取 sessionmaker 对象
        session_maker = get_db_session()
        # 创建会话实例
        session = session_maker()

        # 使用 text 函数明确声明 SQL 查询语句为文本表达式
        result = session.execute(text(f"SELECT * FROM {table_name}"))
        rows = result.fetchall()

        # 获取列名并转换为列表
        column_names = list(result.keys())

        # 创建一个新的工作簿和工作表
        workbook = Workbook()
        sheet = workbook.active

        # 写入列名
        sheet.append(column_names)

        # 写入数据行
        for row in rows:
            sheet.append(row)

        # 保存工作簿为 XLSX 文件
        workbook.save(output_file)
        print(f"数据已成功导出到 {output_file}")

    except Exception as e:
        print(f"发生错误: {e}")


if __name__ == "__main__":
    # 请根据实际情况修改以下数据库连接信息
    table_name = "virtualscene"
    output_file = "output.xlsx"
    export_to_xlsx(table_name, output_file)